import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import math
import scipy.constants as C
import Init


def calc_surface(XData, YData, nx, ny, minx=None, maxx=None, miny=None, maxy=None, T=298.15):
    minx = minx or min(XData)
    maxx = maxx or max(XData)
    miny = miny or min(YData)
    maxy = maxy or max(YData)
    Ix = maxx - minx
    Iy = maxy - miny
    kB = C.k
    An = C.Avogadro

    # 分块统计
    hist, xedges, yedges = np.histogram2d(XData, YData, bins=[nx, ny],
                                          range=np.array([[minx, maxx], [miny, maxy]]))
    # 计算玻尔兹曼函数
    LnPmax = math.log(np.max(hist))
    hist[hist == 0] = -1
    hist[hist > 0] = -0.001 * kB * An * T * (np.log(hist[hist > 0]) - LnPmax)
    hist[hist == -1] = np.max(hist[hist > 0])
    return hist.T


def plot_contour(ax, fig, surf, minx, maxx, miny, maxy, xname='axis1', yname='axis2'):
    ny, nx = surf.shape
    x = np.linspace(minx, maxx, nx)
    y = np.linspace(miny, maxy, ny)

    X, Y = np.meshgrid(x, y)
    cs = ax.contourf(X, Y, surf, cmap='rainbow', levels=100)
    fig.colorbar(cs)
    ax.set_xlabel(xname)
    ax.set_ylabel(yname)


if __name__ == '__main__':
    rg_table = pd.read_csv('Rg.csv')['Rg']
    hbnum_table = pd.read_csv('hbnum.csv')['hbnum']

    # params
    minx, maxx, miny, maxy = 1.2, 1.6, 10, 50
    nx, ny = 16, 16
    xname, yname = 'Rg (nm)', 'Number of Hydrogen Bonds'

    fig, axs = Init.init_figs(1, 1)

    surf = calc_surface(rg, hbnum, nx, ny, minx, maxx, miny, maxy)
    plot_contour(axs, surf, minx, maxx, miny, maxy, xname, yname)

    plt.show(f)
